'''
使用X-AnyLabeling 标注并导出数据文件，图像和标签分别在 data 目录和 classes.txt 文件中。
脚本和文件放到同一目录：
>ls
classes.txt data split.py
'''

import os
import shutil
from pathlib import Path
import random
import yaml   # 如未安装：pip install pyyaml

# ------------- 1. 路径与类别 -------------
data_dir   = Path("./data")
cls_file   = Path("./classes.txt")
out_root   = Path("./")                 # 最终 images / labels 会放在这里
images_dir = out_root / "images"
labels_dir = out_root / "labels"

# 读取类别
with cls_file.open("r", encoding="utf-8") as f:
    class_names = [line.strip() for line in f if line.strip()]
nc = len(class_names)

# 创建输出目录
for d in (images_dir, labels_dir):
    d.mkdir(parents=True, exist_ok=True)

# ------------- 2. 收集并打乱数据 -------------
files   = sorted(data_dir.glob("*.jpg")) + sorted(data_dir.glob("*.jpeg"))
labels  = sorted(data_dir.glob("*.txt"))
assert len(files) == len(labels), "图片和标签数量不一致！"

data = list(zip(files, labels))
random.seed(42)
random.shuffle(data)

# ------------- 3. 划分比例 -------------
total       = len(data)
train_split = int(total * 0.7)
val_split   = int(total * 0.8)

train_data = data[:train_split]
val_data   = data[train_split:val_split]
test_data  = data[val_split:]

# ------------- 4. 拷贝文件 + 生成 txt -------------
def save_files(data_list, subset):
    img_dir = images_dir / subset
    lbl_dir = labels_dir / subset
    img_dir.mkdir(parents=True, exist_ok=True)
    lbl_dir.mkdir(parents=True, exist_ok=True)

    list_txt = out_root / f"{subset}.txt"
    with list_txt.open("w", encoding="utf-8") as f:
        for img, lbl in data_list:
            shutil.copy(img, img_dir / img.name)
            shutil.copy(lbl, lbl_dir / lbl.name)
            # 写相对路径，方便迁移
            f.write(f"{img_dir / img.name}\n")

save_files(train_data, "train")
save_files(val_data,   "val")
save_files(test_data,  "test")

# ------------- 5. 生成 dataset.yaml -------------
yaml_path = out_root / "dataset.yaml"
yaml_data = {
    "path": str(out_root.resolve()),  # 也可以写相对 "./"
    "train": "images/train",
    "val":   "images/val",
    "test":  "images/test",
    "nc": nc,
    "names": class_names
}
with yaml_path.open("w", encoding="utf-8") as f:
    yaml.dump(yaml_data, f, allow_unicode=True, default_flow_style=False)

print("✅ 数据集划分完成，dataset.yaml 已生成！")